package cn.edu.flink.tutorial.watermark;

import cn.edu.flink.tutorial.source.UDSourceFunction;
import org.apache.flink.api.common.eventtime.*;
import org.apache.flink.api.common.functions.ReduceFunction;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.windowing.ProcessWindowFunction;
import org.apache.flink.streaming.api.windowing.assigners.TumblingEventTimeWindows;
import org.apache.flink.streaming.api.windowing.time.Time;
import org.apache.flink.util.Collector;

import java.time.Duration;


public class Strategy {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);

        DataStream<Tuple2<Long, Long>> streamSource = env.addSource(new UDSourceFunction()).map(x -> Tuple2.of(x, 1L)).returns(Types.TUPLE(Types.LONG, Types.LONG));

        DataStream<Tuple2<Long, Long>> watermarkedStream = generator(streamSource);

        window(watermarkedStream);
        env.execute();
    }

    public static DataStream<Tuple2<Long, Long>> generator(DataStream<Tuple2<Long, Long>> streamSource) {
        SingleOutputStreamOperator<Tuple2<Long, Long>> watermarkedStream = streamSource
                .assignTimestampsAndWatermarks(WatermarkStrategy
                        .forGenerator(new WatermarkGeneratorSupplier<Tuple2<Long, Long>>() {
                            @Override
                            public WatermarkGenerator<Tuple2<Long, Long>> createWatermarkGenerator(Context context) {
                                return new WatermarkGenerator<Tuple2<Long, Long>>() {
                                    private long maxTimestamp;
                                    private final long outOfOrdernessMillis = 5000;

                                    @Override
                                    public void onEvent(Tuple2<Long, Long> event, long eventTimestamp, WatermarkOutput output) {
                                        maxTimestamp = Math.max(maxTimestamp, eventTimestamp);
                                    }

                                    @Override
                                    public void onPeriodicEmit(WatermarkOutput output) {
                                        // 周期调用 通过 env.getConfig().setAutoWatermarkInterval() 设置 默认 200ms
                                        System.out.println(maxTimestamp);
                                        output.emitWatermark(new Watermark(maxTimestamp - outOfOrdernessMillis - 1));
                                    }
                                };
                            }
                        })
                        .withTimestampAssigner(new SerializableTimestampAssigner<Tuple2<Long, Long>>() {
                            @Override
                            public long extractTimestamp(Tuple2<Long, Long> element, long recordTimestamp) {
                                return element.f0;
                            }
                        }));

        return watermarkedStream;
    }

    public static DataStream<Tuple2<Long, Long>> newWatermarkStrategy(DataStream<Tuple2<Long, Long>> streamSource) {
        SingleOutputStreamOperator<Tuple2<Long, Long>> watermarkedStream = streamSource
                .assignTimestampsAndWatermarks(new WatermarkStrategy<Tuple2<Long, Long>>() {
                    @Override
                    public TimestampAssigner<Tuple2<Long, Long>> createTimestampAssigner(TimestampAssignerSupplier.Context context) {
                        return new SerializableTimestampAssigner<Tuple2<Long, Long>>() {
                            @Override
                            public long extractTimestamp(Tuple2<Long, Long> element, long recordTimestamp) {
                                return element.f0;
                            }
                        };
                    }

                    @Override
                    public WatermarkGenerator<Tuple2<Long, Long>> createWatermarkGenerator(WatermarkGeneratorSupplier.Context context) {
                        return new WatermarkGenerator<Tuple2<Long, Long>>() {
                            private long maxTimestamp;
                            private final long outOfOrdernessMillis = 5000;

                            @Override
                            public void onEvent(Tuple2<Long, Long> event, long eventTimestamp, WatermarkOutput output) {
                                maxTimestamp = Math.max(maxTimestamp, eventTimestamp);
                            }

                            @Override
                            public void onPeriodicEmit(WatermarkOutput output) {
                                // 周期调用 通过 env.getConfig().setAutoWatermarkInterval() 设置 默认 200ms
                                System.out.println(maxTimestamp);
                                output.emitWatermark(new Watermark(maxTimestamp - outOfOrdernessMillis - 1));
                            }
                        };
                    }
                });

        return watermarkedStream;
    }

    public static DataStream<Tuple2<Long, Long>> boundedOutOfOrderness(DataStream<Tuple2<Long, Long>> streamSource) {
        SingleOutputStreamOperator<Tuple2<Long, Long>> watermarkedStream = streamSource
                .assignTimestampsAndWatermarks(WatermarkStrategy
                        .<Tuple2<Long, Long>>forBoundedOutOfOrderness(Duration.ofSeconds(5))
                        .withTimestampAssigner(new SerializableTimestampAssigner<Tuple2<Long, Long>>() {
                            @Override
                            public long extractTimestamp(Tuple2<Long, Long> element, long recordTimestamp) {
                                return element.f0;
                            }
                        }));

        return watermarkedStream;
    }

    public static DataStream<Tuple2<Long, Long>> monotonousTimestamps(DataStream<Tuple2<Long, Long>> streamSource) {
        SingleOutputStreamOperator<Tuple2<Long, Long>> watermarkedStream = streamSource
                .assignTimestampsAndWatermarks(WatermarkStrategy
                        .<Tuple2<Long, Long>>forMonotonousTimestamps()
                        .withTimestampAssigner(new SerializableTimestampAssigner<Tuple2<Long, Long>>() {
                            @Override
                            public long extractTimestamp(Tuple2<Long, Long> element, long recordTimestamp) {
                                return element.f0;
                            }
                        }));
        return watermarkedStream;
    }

    public static void window(DataStream<Tuple2<Long, Long>> watermarkedStream) {
        watermarkedStream
                .keyBy(x -> 1)
                .window(TumblingEventTimeWindows.of(Time.seconds(5)))
                .reduce(
                        new ReduceFunction<Tuple2<Long, Long>>() {
                            @Override
                            public Tuple2<Long, Long> reduce(Tuple2<Long, Long> value1, Tuple2<Long, Long> value2) throws Exception {
                                value1.f1 = value2.f0;
                                return value1;
                            }
                        },
                        new ProcessWindowFunction<Tuple2<Long, Long>, Object, Integer, org.apache.flink.streaming.api.windowing.windows.TimeWindow>() {
                            @Override
                            public void process(Integer integer, Context context, Iterable<Tuple2<Long, Long>> elements, Collector<Object> out) throws Exception {
                                System.out.println("key:" + integer);
                                System.out.println("currentProcessingTime:" + context.currentProcessingTime());
                                System.out.println("currentWatermark:" + context.currentWatermark());
                                System.out.println("dely:" + (context.currentProcessingTime() - context.currentWatermark()));
                                System.out.println("window:" + context.window());
                                System.out.println("elements:" + elements);
                                System.out.println();
                            }
                        });
    }
}
